#include <iostream>
#include <fstream>
#include <iterator>
#include <vector>
#include <map>
#include <string>
#include <sstream>
#include "QtCore/QFile"
#include "QtCore/QDir"
#include "QtCore/QList"
#include "convert_utils.h"
#include "main_widget_c.h"

namespace ConvertUtils
{
    using std::map;
    using std::string;
    using std::vector;

    class IMGAnnotation
    {
    private:
        string get_yolov5_label(const vector<string> &boxes, map<string, int> &class_to_idx)
        {
            // 获取索引
            if (class_to_idx.count(boxes[0]) == 0)
            {
                class_to_idx.insert(std::make_pair(boxes[0], class_to_idx.size()));
            }
            int idx = class_to_idx[boxes[0]];
            // 转换为yolov5格式
            string out = std::to_string(idx);
            float box1[4];
            for (int i = 0; i < 4; i++)
            {
                box1[i] = std::stof(boxes[i + 1]);
            }
            float box2[4] = {
                (box1[0] + box1[2]) / 2 / width,
                (box1[1] + box1[3]) / 2 / height,
                (box1[2] - box1[0]) / width,
                (box1[3] - box1[1]) / height};
            for (int i = 0; i < 4; i++)
            {
                out += " " + std::to_string(box2[i]);
            }
            return out;
        }

    public:
        string filename;
        int width;
        int height;
        vector<vector<string>> objects;
        IMGAnnotation(string output_path, string w, string h) : filename(output_path),
                                                                width(std::stoi(w)), height(std::stoi(h))
        {
        }

        void append_object(vector<string> &obj)
        {
            objects.push_back(obj);
        }

        void save(const string &path, map<string, int> &class_to_idx)
        {
            std::ofstream file;
            file.open(path);
            if (!file.is_open())
            {
                return;
            }
            if (objects.size() == 0)
            {
                // warning
                return;
            }
            for (auto &obj : objects)
            {
                file << get_yolov5_label(obj, class_to_idx) << std::endl;
            }
            file.close();
        }
    };

    int parse_a_tags(const string &buffer, const string &tag_name, string &content, int index)
    {
        int size = buffer.size();
        string cur_tag;
        if (index >= size)
            return -1;
        // 定位标到 <
        while (index + 1 < size and (not(buffer[index] == '<' and buffer[index + 1] != '/')))
        {
            index++;
        }

        if (index + 1 >= size)
        {
            return -1;
        }
        // 读取到标签名
        index++;
        while (buffer[index] != '>')
        {
            cur_tag.push_back(buffer[index]);
            index++;
        }

        index++;
        // 不匹配的标签名
        if (cur_tag != tag_name)
        {
            return parse_a_tags(buffer, tag_name, content, index);
        }
        else
        {
            while (index < size and buffer[index] != '<')
            {
                content.push_back(buffer[index]);
                index++;
            }
        }
        return index;
    }

    void create_yaml_and_commands(const QDir& output_dir, const map<string, int> &class_to_idx){
        std::ofstream file;
        file.open(output_dir.filePath("kaggle_data.yaml").toStdString());
        if (!file.is_open())
        {
            std::cout << "open file failed" << std::endl;
            return;
        }
        file << "names:" << std::endl;
        for(auto item : class_to_idx){
            file << "  " << item.second << ": " << item.first << std::endl;
        }
        QString dir_name = output_dir.dirName();
        file << "path: /kaggle/input/" << dir_name.toStdString() << std::endl;
        file << "train: images/train" << std::endl;
        file << "val: images/val" << std::endl;
        file.close();

        file.open(output_dir.filePath("commands.txt").toStdString());
        file << "!git clone https://github.com/ultralytics/yolov5.git && wandb off" << std::endl;
        file << "!pip install ultralytics" << std::endl;
        file << "!python yolov5/train.py --img 640 --epochs 10 --data /kaggle/input/" 
             <<dir_name.toStdString()
             <<"/kaggle_data.yaml --weights yolov5s.pt"
             << std::endl;
        file.close();
    }

    int convert_xml(const string &xml_path, const string &output_path, map<string, int> &class_to_idx)
    {
        std::ifstream file;
        file.open(xml_path);
        if (!file.is_open())
        {
            std::cout << "open file failed" << std::endl;
            return -1;
        }
        std::istream_iterator<char> begin(file);
        std::istream_iterator<char> end;
        string buffer(begin, end);
        file.close();

        string name;
        string width;
        string height;
        int index = 0;
        index = parse_a_tags(buffer, "filename", name, index);
        index = parse_a_tags(buffer, "width", width, index);
        index = parse_a_tags(buffer, "height", height, index);
        IMGAnnotation img(name, width, height);
        while (true)
        {
            string name, xmin, ymin, xmax, ymax;
            index = parse_a_tags(buffer, "name", name, index);
            if (index == -1)
                break;
            index = parse_a_tags(buffer, "xmin", xmin, index);
            index = parse_a_tags(buffer, "ymin", ymin, index);
            index = parse_a_tags(buffer, "xmax", xmax, index);
            index = parse_a_tags(buffer, "ymax", ymax, index);
            vector<string> obj = {name, xmin, ymin, xmax, ymax};
            img.append_object(obj);
        }
        if (img.objects.size() == 0)
        {
            return -1;
        }
        img.save(output_path, class_to_idx);
        return 0;
    }

    int rfind(const QString &s)
    {
        for (int i = s.size() - 1; i >= 0; i--)
        {
            if (s[i] == '\\' || s[i] == '/')
            {
                return -1;
            }
            if (s[i] == '.')
            {
                return i;
            }
        }
        return -1;
    }

    QString basename(const QString &s)
    {
        int index = rfind(s);
        if (index == -1)
        {
            return s;
        }
        return s.left(index);
    }
    

    void convert(
        const QStringList& img_dirs,
        float train_rate,
        const QString& output_dir,
        Window *win)
    {
        QStringList warning;
        map<string, int> class_to_idx;
        QDir qoutput_dir(output_dir);
        QDir train_img_dir = qoutput_dir.filePath("images/train");
        QDir train_label_dir = qoutput_dir.filePath("labels/train");
        QDir val_img_dir = qoutput_dir.filePath("images/val");
        QDir val_label_dir = qoutput_dir.filePath("labels/val");
        for(const auto& dir : {train_img_dir, train_label_dir, val_img_dir, val_label_dir}){
            if(!dir.exists()){
                dir.mkpath(dir.absolutePath());
            }
        }

        int img_count = 0;
        for(const auto& img_dir : img_dirs){
            QDir qdir(img_dir);
            
            QStringList img_list = qdir.entryList(QDir::Files);
            img_list = img_list.filter(".jpg");
            img_count += img_list.size();

        }


        int index = 0;
        int train_idx = 0;
        int val_idx = 0;

        int r = 1 / (1-train_rate);

        


        for(const auto& img_dir : img_dirs){
            QDir qimg_dir(img_dir);
            QDir qxml = qimg_dir;
            qxml.cdUp();
            qxml.cd("Annotations");

            // 默认采用的是名称排序而不是自然排序
            QStringList img_list = qimg_dir.entryList(QDir::Files);

            for (const auto &img : img_list)
            {   

                QString file_basename = ConvertUtils::basename(img);
                const QDir& img_out_idr = index % r == 0 ? val_img_dir : train_img_dir;
                const QDir& label_out_idr = index % r == 0 ? val_label_dir : train_label_dir;
                int file_id = index % r == 0 ? val_idx : train_idx;

                // 创建label
                // 源xml路径
                QString xml_source_path = qxml.filePath(file_basename + ".xml");
                // 目的xml路径
                QString xml_target_path = label_out_idr.filePath(QString::number(file_id) + ".txt");

                // std::cout << xml_source_path.toStdString() << std::endl;
                // std::cout << xml_target_path.toStdString() << std::endl;
                int ret = convert_xml(xml_source_path.toStdString(), xml_target_path.toStdString(), class_to_idx);

                if (ret == -1){
                    warning.push_back(xml_source_path + " 存在问题，已跳过");
                    continue;
                }
                // 移动文件
                // 源图片路径
                QString img_source_path = qimg_dir.filePath(img);
                // 目的图片路径
                QString img_target_path = img_out_idr.filePath(QString::number(file_id) + ".jpg");
                // 移动图片
                // std::cout << img_source_path.toStdString() << std::endl;
                // std::cout << img_target_path.toStdString() << std::endl;
                QFile::copy(img_source_path, img_target_path);


                emit win->log_signal(img_target_path);
                emit win->progress_signal((int)(index * 100.0 / img_count));

                if (index % r == 0)
                {
                    val_idx += 1;
                }
                else
                {
                    train_idx += 1;
                }
                index += 1;
                // 移动图片
                // img_file.copy(img_out_idr + "/" + img_
            }

            for(int i=0; i<warning.size(); i++){
                emit win->log_signal(warning[i]);
            }

            create_yaml_and_commands(qoutput_dir, class_to_idx);
            emit win->progress_signal(100);  // 强制设置为100
        }
    }



    QString join(QString path1, QString path2)
    {
        if (path1.back() == '\\' || path1.back() == '/')
        {
            path1 = path1.remove(path1.size() - 1, 1);
        }
        path1 += path2;
        return path1;
    }
    //    QString parent(QString)
}